data_dir = '/media/dengyunfei/6T/data/datasets/supremezxc/nlpcc_data_origin.json'
model_dir = '/media/dengyunfei/6T/data/models/huggingface/mengzi-t5-base'
save_dir = '/media/dengyunfei/6T/data/logs/abstract_model'
# 导入相关包
# import torch
# from datasets import Dataset
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,T5ForConditionalGeneration,DataCollatorForSeq2Seq,Seq2SeqTrainer,Seq2SeqTrainingArguments
# 加载数据集
datasets = load_dataset(path='json', data_files=data_dir,split='train')
"""
path 
1、 可以指定文件的所在目录，如果里面存在一个和上级目录同名的脚本。例如：
dataset = load_dataset("logistic_data", split="train",trust_remote_code=True)
print(dataset[0:40]) logistic_data 目录下有一个自定义的处理脚本 名称为 logistic_data.py 数据将按照指定脚本来执行并处理。
注意，trust_remote_code=True 这个设定才可以使用自定义脚本进行处理。

2、可以直接指定处理数据的脚本。例如：
dataset = load_dataset("dataset_handler.py", split="train",trust_remote_code=True)
print(dataset[0:40])

3、默认指定文件所在目录也可以，目录中没有处理脚本的，将按照默认识别的方式来处理
datasets = load_dataset(path="/media/dengyunfei/6T/data/datasets/supremezxc/", split='train')
datasets = load_dataset(path="/media/dengyunfei/6T/data/datasets/supremezxc/", data_files=data_dir,split='train')

4、直接指定处理方式 需要补充指定 data_files
datasets = load_dataset(path=“json", data_files=data_dir,split='train')
"""
datasets = load_dataset(path="json", data_files=data_dir,split='train')
# print(type(datasets)) datasets.arrow_dataset.Dataset

"""
切分数据集合 test_size 如果是一个小于1的小数，小数的值代表了测试集合的比例。
如果大于1的整数，则代表了测试集合保留多少个数据。
"""
# datasets = datasets.train_test_split(test_size=0.90)
# datasets = datasets["train"]
datasets = datasets.train_test_split(0.1,seed=42)
print(datasets)
# print(datasets["train"])
# print(datasets["test"])
print(datasets["train"][0])
# 只有 content和title

# 数据处理
tokenizer = AutoTokenizer.from_pretrained(model_dir)

def process_func(examples):
    contents = ["摘要生成:\n" + e for e in examples["content"]]
    # 分开处理
    inputs =  tokenizer(contents, max_length=1000000, truncation=True)
    labels =  tokenizer(text_target=examples["title"], max_length=128, truncation=True)
    inputs["labels"] =  labels["input_ids"]
    return inputs

tokenizer_ds =  datasets.map(process_func,batched=True)
#
# # 将数据处理成如下这个样子
# print(tokenizer.decode(tokenizer_ds["train"][0]["input_ids"]))
# print(tokenizer.decode(tokenizer_ds["train"][0]["labels"]))
#
# # 创建模型
# model =  AutoModelForSeq2SeqLM.from_pretrained(model_dir)
#
# # 创建评估函数
# import numpy as np
# from rouge_chinese import Rouge
# rouge = Rouge()
#
# def compute_metric(eval_preds):
#     predictions,labels = eval_preds
#     decoded_preds = tokenizer.batch_decode(predictions,skip_special_tokens=True)
#     # labels 有-100 因此需要进行填充操作,因为labels有-100 不属于特殊字符，因此需要填充为特殊字符
#     labels = np.where(labels!=-100,labels,tokenizer.pad_token_id)
#     decoded_labels = tokenizer.batch_decode(labels,skip_special_tokens=True)
#     # 基于中文的字来做评估
#     decoded_preds = [" ".join(pred.strip()) for pred in decoded_preds]
#     decoded_labels = [" ".join(label.strip()) for label in decoded_labels]
#     # 该方法计算 Rouge，标准的文字评估方案
#     scores = rouge.get_scores(decoded_preds,decoded_labels,avg=True)
#     return {
#         "rouge-1":scores["rouge-1"]["f"],
#         "rouge-2":scores["rouge-2"]["f"],
#         "rouge-l":scores["rouge-l"]["f"]
#     }
#
# args = Seq2SeqTrainingArguments(
#     output_dir=save_dir,
#     per_device_train_batch_size=4,
#     per_device_eval_batch_size=4,
#     gradient_accumulation_steps=2,
#     logging_steps = 100,
#     num_train_epochs=4,
#     eval_strategy="epoch",
#     save_strategy="epoch",
#     metric_for_best_model="rouge-l",
#     # predict_with_generate 这个参数必须设置为True，否则无法做评估
#     predict_with_generate = True
# )
# args.num_train_epochs              # total number of training epochs
#
# trainer =  Seq2SeqTrainer(
#     args = args,
#     model =model,
#     train_dataset = tokenizer_ds["train"],
#     eval_dataset = tokenizer_ds["test"],
#     compute_metrics = compute_metric,
#     tokenizer =tokenizer,
#     data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer)
#
# )
#
# # 开始训练
# trainer.train()
#
# # 模型推理
# # 导入相关包
# import torch
# import os
# from datasets import Dataset,load_dataset
# from transformers import pipeline
# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,T5ForConditionalGeneration,DataCollatorForSeq2Seq,Seq2SeqTrainer,Seq2SeqTrainingArguments
#
# # data_dir = '/data/datasets/supremezxc/nlpcc_data.json'
# model_dir = '/data/models/huggingface/mengzi-t5-base'
# save_dir = '/data/logs/abstract_model'
# # E:\data\logs\abstract_model\checkpoint-4900
# tokenizer = AutoTokenizer.from_pretrained(model_dir)
# model = AutoModelForSeq2SeqLM.from_pretrained(os.path.join(save_dir,'checkpoint-4900')).to('cuda')
# pipe = pipeline("text2text-generation",model=model,tokenizer=tokenizer,device=0)
#
# pipe("摘要生成:\n"+datasets["test"][-2]["content"],max_length=128,do_sample=True)
#
# text = '10月1日晚8点20，WTT中国大满贯男单16强开赛，王楚钦对阵丹麦选手安德斯·林德。这是王楚钦本次大赛的重要一战，他在之前的预选赛中已成功晋级。比赛首局，王楚钦11-13安德斯·林德，比分0:1落后；第二句，王楚钦扳回一局；第三局，王楚钦6-11安德斯·林德，比分1-2落后；第四局林德上场后连续进攻取得领先优势。随后被王楚钦反超比分。然而关键时刻王楚钦自己也出现打丢的情况，叫出暂停也未能调整好心态，最终以7-11无缘晋级男单十六强。'
#
# pipe("摘要生成:\n"+text,max_length=128,do_sample=True)